import copy

import numpy as np


class CliffWalkingEnv:
    """
    悬崖漫步环境
        这段代码实现了：
        状态空间：48 个格子
        动作空间：4 个方向
        奖励机制：普通格子：-1
        掉下悬崖：-100（并结束）
        到达终点：0（并结束）
    返回:
        得出每个格子的四个方向的转移函数[(p, next_state, reward, done)]
        得出所有的各自的四个方向的转移矩 P
    """
    def __init__(self, ncol=12, nrow=4):
        self.ncol = ncol  # 定义网格世界的列
        self.nrow = nrow  # 定义网格世界的行
        # 转移矩阵P[state][action] = [(p, next_state, reward, done)]包含下一个状态和奖励
        self.P = self.createP()
        print(f"状态转移矩阵: {self.P}")

    def createP(self):
        # 初始化, 每个状态都有 4 个动作
        P = [[[] for j in range(4)] for i in range(self.nrow * self.ncol)]
        # 4种动作, change[0]:上,change[1]:下, change[2]:左, change[3]:右。坐标系原点(0,0)
        # 坐标系定义在左上角
        change = [[0, -1], [0, 1], [-1, 0], [1, 0]]
        for i in range(self.nrow):  # 遍历行
            for j in range(self.ncol):  # 遍历列
                for a in range(4):  # 遍历动作
                    # 位置在悬崖或者目标状态,因为无法继续交互,任何动作奖励都为0
                    if i == self.nrow - 1 and j > 0:    # 定义第三行(也就是最后一行), (3,1) ~ (3,11) → 全部都被当成“终止状态”，包含 悬崖 和 终点。
                        P[i * self.ncol + j][a] = [(1, i * self.ncol + j, 0, True)]  # i * self.ncol + j表示当前的位置
                        """
                        P[47][0] = [(1, 47, 0, True)]
                        P[47][1] = [(1, 47, 0, True)]
                        P[47][2] = [(1, 47, 0, True)]
                        P[47][3] = [(1, 47, 0, True)]
                        """
                        continue
                    # 其他位置, next_x, next_y就表示下一个位置的x,y的坐标
                    next_x = min(self.ncol - 1, max(0, j + change[a][0]))   # change[a][0]表示向左右变化
                    next_y = min(self.nrow - 1, max(0, i + change[a][1]))   # change[a][1]表示向上下变化
                    next_state = next_y * self.ncol + next_x    # 计算下一个状态的一维坐标
                    reward = -1
                    done = False
                    # 下一个位置在悬崖或者终点
                    if next_y == self.nrow - 1 and next_x > 0:
                        done = True
                        if next_x != self.ncol - 1:  # 下一个位置在悬崖
                            reward = -100
                    P[i * self.ncol + j][a] = [(1, next_state, reward, done)]
        # P是状态转移函数, 会将奖励和环境结合起来组合成四元组(转移概率、下一个状态、奖励、是否终止)
        return P


class PolicyIteration:
    """ 策略迭代算法 """
    def __init__(self, env, theta, gamma):
        self.env = env
        self.v = [0] * self.env.ncol * self.env.nrow  # 初始化价值为0
        self.pi = [[0.25, 0.25, 0.25, 0.25] for i in range(self.env.ncol * self.env.nrow)]  # 初始化为均匀随机策略, 每一个动作的价值初始化为0.25
        """
        策略的概率Π(a|s): 因为没有学到任何东西, 总共有四个方向, 所有每个方向的动作概率为0.25
        self.pi = [
            [0.25, 0.25, 0.25, 0.25], 动作0(下)的概率为0.25, 动作1(上)的概率为0.25, 动作2(左)的概率为0.25, 动作3(右)的概率为0.25,
            [0.25, 0.25, 0.25, 0.25],
            ...,
            [0.25, 0.25, 0.25, 0.25]
        ]
        """
        self.theta = theta  # 策略评估收敛阈值
        self.gamma = gamma  # 折扣因子

    def policy_evaluation(self):  # 策略评估
        cnt = 1  # 计数器
        while 1:
            max_diff = 0    # 状态价值函数的更新幅度里最大的那个差值
            new_v = [0] * self.env.ncol * self.env.nrow     # 新价值=[0]列表×列×行  也就是得到一个列×行的全零一位列表[0,0,...,0]
            for s in range(self.env.ncol * self.env.nrow):  # 循环取12×4=48的s值, 0,1,2...,47
                qsa_list = []  # 开始计算状态s下的所有Q(s,a)价值
                for a in range(4):  # 遍历四个动作
                    qsa = 0  # 初始化动作价值
                    for res in self.env.P[s][a]:    # 遍历当前格子的当前动作的状态转移矩阵, 这里其实只有一个转移函数
                        p, next_state, r, done = res
                        # 贝尔曼方程的一般形式(s, a, s'), 所以环境转移概率P(s'|s,a)=1
                        qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))   # qsa计算出来的是当前格子的当前动作价值, 也就是这个格子的某个动作价值
                    qsa_list.append(self.pi[s][a] * qsa)    # self.pi[s][a] * qsa: 计算当前策略下的, 选中的概率×动作的价值; 用来计算价值函数
                new_v[s] = sum(qsa_list)  # 状态价值函数和动作价值函数之间的关系: 状态价值=每个动作价值按照该动作被选择的概率做加权平均
                # abs(): 取绝对值
                # abs(new_v[s] - self.v[s]): 状态的新值和旧值相差的绝对值是多少
                # max(max_diff, abs(new_v[s] - self.v[s])): 取当前的状态差值和max_diff, 哪个大取哪个, 循环更新完毕后, max_diff就是最大的状态价值函数的最大值
                max_diff = max(max_diff, abs(new_v[s] - self.v[s]))
            self.v = new_v
            if max_diff < self.theta:
                break  # 满足收敛条件,退出评估迭代
            cnt += 1
        print("策略评估进行%d轮后完成, 计算出所有动作的状态价值如下" % cnt)

    def policy_improvement(self):  # 策略提升
        for s in range(self.env.nrow * self.env.ncol):  # 循环取12×4=48的s值, 0,1,2...,47
            qsa_list = []
            for a in range(4):  # 遍历四个动作
                qsa = 0
                for res in self.env.P[s][a]:    # 遍历当前格子的当前动作的状态转移矩阵, 这里其实只有一个转移函数
                    p, next_state, r, done = res
                    qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))
                qsa_list.append(qsa)
            # Q 值衡量动作的未来价值，最大 Q 值代表当前状态下最优的选择；所以策略提升就是“往最大 Q 值的方向走”。
            maxq = max(qsa_list)    # 当前格子的所有动作价值函数中最大的值;
            cntq = qsa_list.count(maxq)  # 计算有几个动作得到了最大的Q值
            # 让这些动作均分概率
            # self.pi[s] = [1 / cntq if q == maxq else 0 for q in qsa_list]
            new_probs = []
            for q in qsa_list:  # 遍历当前格子的四个动作价值
                if q == maxq:  # 如果这个动作的价值等于最大动作价值
                    new_probs.append(1 / cntq)  # 给它分配 1/cntq 的概率
                else:
                    new_probs.append(0)  # 否则概率为0
            # 打印这个过程
            print(f"状态 {s:2d}: Q值 = {qsa_list}, 最大Q = {maxq:.3f}, "
                  f"最优动作数量 = {cntq}, 新策略 = {new_probs}")
            self.pi[s] = new_probs
        print("策略提升完成")
        return self.pi

    def policy_iteration(self, action_meaning, disaster, end):  # 策略迭代
        iteration = 1
        while 1:    # 无限循环
            self.print_value_policy(iteration, action_meaning, disaster, end)
            self.policy_evaluation()    # 策略评估: 给定当前策略 π，使用贝尔曼方程计算出每个状态的价值函数 V(s)。
            old_pi = copy.deepcopy(self.pi)  # 将列表进行深拷贝(把当前策略保存一份),方便接下来进行比较
            new_pi = self.policy_improvement()  # 策略提升: 在评估的价值函数 V(s) 基础上，改进策略 π。找到让价值更高的动作，把概率分配给这些最优动作。得到新策略 new_pi。
            if old_pi == new_pi:
                print("新旧策略完全相同, 策略已经收敛")
                break   # 如果新旧策略完全相同 → 表明策略已经收敛（找不到更优策略了），退出循环。
            iteration += 1

    def print_value_policy(self, cnt, action_meaning, disaster, end):
        print(f"============ 第 {cnt} 次迭代 ============")
        print_agent(self, action_meaning, disaster=disaster, end=end)


class ValueIteration:
    """ 价值迭代算法 """
    def __init__(self, env, theta, gamma):
        self.env = env
        self.v = [0] * self.env.ncol * self.env.nrow  # 初始化价值为0
        self.theta = theta  # 价值收敛阈值
        self.gamma = gamma
        # 价值迭代结束后得到的策略
        self.pi = [None for i in range(self.env.ncol * self.env.nrow)]

    def value_iteration(self):
        cnt = 0
        while 1:
            max_diff = 0
            new_v = [0] * self.env.ncol * self.env.nrow
            for s in range(self.env.ncol * self.env.nrow):
                qsa_list = []  # 开始计算状态s下的所有Q(s,a)价值
                for a in range(4):
                    qsa = 0
                    for res in self.env.P[s][a]:
                        p, next_state, r, done = res
                        qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))
                    qsa_list.append(qsa)  # 这一行和下一行代码是价值迭代和策略迭代的主要区别
                new_v[s] = max(qsa_list)
                max_diff = max(max_diff, abs(new_v[s] - self.v[s]))
            self.v = new_v
            if max_diff < self.theta: break  # 满足收敛条件,退出评估迭代
            cnt += 1
        print("价值迭代一共进行%d轮" % cnt)
        self.get_policy()

    def get_policy(self):  # 根据价值函数导出一个贪婪策略
        for s in range(self.env.nrow * self.env.ncol):
            qsa_list = []
            for a in range(4):
                qsa = 0
                for res in self.env.P[s][a]:
                    p, next_state, r, done = res
                    qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))
                qsa_list.append(qsa)
            maxq = max(qsa_list)
            cntq = qsa_list.count(maxq)  # 计算有几个动作得到了最大的Q值
            # 让这些动作均分概率
            self.pi[s] = [1 / cntq if q == maxq else 0 for q in qsa_list]


def print_agent(agent, action_meaning, disaster=[], end=[]):
    print("状态价值：")
    for i in range(agent.env.nrow):
        for j in range(agent.env.ncol):
            # 为了输出美观,保持输出6个字符
            print('%6.6s' % ('%.3f' % agent.v[i * agent.env.ncol + j]), end=' ')
        print()

    print("策略：")
    for i in range(agent.env.nrow):
        for j in range(agent.env.ncol):
            # 一些特殊的状态,例如悬崖漫步中的悬崖
            if (i * agent.env.ncol + j) in disaster:
                print('****', end=' ')
            elif (i * agent.env.ncol + j) in end:  # 目标状态
                print('EEEE', end=' ')
            else:
                a = agent.pi[i * agent.env.ncol + j]
                pi_str = ''
                for k in range(len(action_meaning)):
                    pi_str += action_meaning[k] if a[k] > 0 else 'o'
                print(pi_str, end=' ')
        print()

def rollout_gridworld(policy, nrow, ncol, start_state, goal_state, fail_states=None, n_episodes=100):
    """
    通用格子环境策略验证

    policy: 策略矩阵
    nrow, ncol: 网格行列数
    start_state: 起点索引
    goal_state: 终点索引
    fail_states: 失败集合 (悬崖/冰洞)，默认为空
    n_episodes: 测试次数
    """
    if fail_states is None:
        fail_states = set()

    success = 0

    for _ in range(n_episodes):
        state = start_state
        steps = 0
        while True:
            action = np.argmax(policy[state])  # 按策略选动作
            row, col = divmod(state, ncol)

            # 动作映射：0=上, 1=下, 2=左, 3=右
            if action == 0 and row > 0: row -= 1
            elif action == 1 and row < nrow - 1: row += 1
            elif action == 2 and col > 0: col -= 1
            elif action == 3 and col < ncol - 1: col += 1

            next_state = row * ncol + col
            steps += 1

            # 判断失败
            if next_state in fail_states:
                break
            # 判断成功
            if next_state == goal_state:
                success += 1
                break

            state = next_state  # 更新状态

    print(f"在 {n_episodes} 次尝试中，成功到达终点 {success} 次，成功率 {success/n_episodes:.2f}")


if __name__ == '__main__':
    env = CliffWalkingEnv()
    theta = 0.001
    gamma = 0.9
    agent = PolicyIteration(env, theta, gamma)  # 初始化策略迭代算法
    action_meaning = ['^', 'v', '<', '>']
    agent.policy_iteration(action_meaning, disaster=list(range(37, 47)), end=[47])    # 策略迭代

    # 尝试机器人走1000次
    nrow, ncol = 4, 12
    start_state = (nrow - 1) * ncol  # 36
    goal_state = nrow * ncol - 1  # 47
    cliff = set(range(start_state + 1, goal_state))  # 37-46

    rollout_gridworld(agent.pi, nrow, ncol, start_state, goal_state, fail_states=cliff, n_episodes=1000)

